from fastapi import APIRouter,Body
from fastapi import FastAPI, WebSocket
from fastapi import FastAPI
from fastapi.responses import JSONResponse, HTMLResponse, FileResponse
from memory.memory import get_Knowledge_Graph
from memory.sql_history import get_latest_messages,insert_message,query_msg_counts,query_kg_message,insert_or_update_kg_message
from llm.llm import llm_ZhipuAI,llm_stream_ZhipuAI
from prompt.role_prompt import get_Miku_prompt
from prompt.prompt_model import get_promptTemplate
from embedding.history_msg_db import find_similar_messages,insert_chat_message
from API.AInews import get_AInews
from API.bing_search import get_bing_search
from zhipuai import ZhipuAI
import requests
from io import BytesIO
from scheduler import get_agent_task
from starlette.websockets import WebSocketDisconnect
from apscheduler.schedulers.background import BackgroundScheduler
from datetime import datetime,timedelta
import uuid
import speech_recognition as sr
from typing import Dict, Any
import asyncio
import websockets
import json
import pymysql.cursors
import re
app = FastAPI()


from fastapi import APIRouter, WebSocket

# Define a custom router with the prefix '/aiagent'
router = APIRouter(prefix='/aiagent')

client = ZhipuAI(api_key="")
# T2A API的配置
group_id = ""
api_key ="eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJHcm91cE5hbWUiOiLojIPlm73mlowiLCJVc2VyTmFtZSI6IuiMg-WbveaWjCIsIkFjY291bnQiOiIiLCJTdWJqZWN0SUQiOiIxNzY4NTM3NjQ1ODMyNDg3MDQzIiwiUGhvbmUiOiIxNzUwMjgzMzAwOSIsIkdyb3VwSUQiOiIxNzY4NTM3NjQ1ODI0MDk4NDM1IiwiUGFnZU5hbWUiOiIiLCJNYWlsIjoiIiwiQ3JlYXRlVGltZSI6IjIwMjQtMDMtMjIgMDA6MDI6MjQiLCJpc3MiOiJtaW5pbWF4In0.a_3fMztYYf2uWaeKLbXQC-6bpiGNMrRznfQ8Na6bTifEkpYXn7-9at9uTmO3NB22D2ehwHanSwzEOpT9eobWp-SgbuCwKCjwzmmrzJci4Cwv9b-g2xfuRK8WUPDSxJfbS7ZFx3l1Bg2_bbPanrlsJrleZHgOXEKQBSl28fN8bS5X7OTfhzSN9mV5VUtv-4XZW5Jwzy5D8MLf5uOp5J1c6nyvGhk2MhGuXQWOuieKIzs4AcRSkzxWxRvPQzFLe7x5dEKZuEPESjsGDWqcCp1p-DuCl43J_V__AZI9cOPcB9EWHeDlA35X8TLX2hnnd6T3xVm0RbpfDtIazHoZ5cSY3w"

tts_url = "https://api.minimax.chat/v1/text_to_speech"
headers = {
    "Authorization": f"Bearer {api_key}",
    "Content-Type": "application/json"
}

def generate_tts(text):
    payload = {
        "text": text,
        "model": "speech-02",  # 选择合适的模型
        "speed": 1.33,
        "vol": 1.0,
        "pitch": 0,
        "timber_weights": [
            {
                "voice_id": "audiobook_male_2",
                "weight": 1
            },
            {
                "voice_id": "audiobook_male_2",
                "weight": 1
            }
        ],
        # 其他参数可以根据需要添加
    }
    response = requests.post(tts_url, headers=headers, json=payload)
    if response.status_code == 200:
        return response.content
    else:
        return None

scheduler = BackgroundScheduler()
scheduler.start()


async def job(websocket, event, user_id):
    response = f"【事件】记得{event}哦\n"
    print(response)
    agent_id = "666666"
    insert_message(user_id, agent_id, "assistant", response)
    scheduler.shutdown()

    content = {"msg": "任务提醒", "data": response, "type": "task", "code": 2000}

    await websocket.send_json(content)



 #www.hspace.top/aiagent/ws
@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            try:
                user_id = body.get("user_id")
                username = body.get("username")
                # data = await websocket.receive_text()
                body = await websocket.receive_json()
                msg_type = body.get("type")  # 假设消息中包含一个"type"字段来指示消息类型

                if msg_type == "text":
                    question = body.get("text")
                    # 处理文本消息
                elif msg_type == "voice":
                    voice_data = body.get("data")  # 假设语音数据作为"data"字段发送
                    # 创建一个Recognizer对象
                    recognizer = sr.Recognizer()
                    # 将Base64编码的语音数据转换为音频文件
                    audio_file = sr.AudioFileData(voice_data)
                    # 使用Recognizer对象处理音频文件
                    with sr.AudioFile(audio_file) as source:
                        audio = recognizer.record(source)
                    # 尝试使用Google Web Speech API进行语音识别
                    try:
                        question = recognizer.recognize_google(audio, language='zh-cn')

                    except sr.UnknownValueError:
                        print("无法理解音频")
                        await websocket.send_json({"msg": "无法理解音频","type":"text", "code": 2001})

                    except sr.RequestError as e:
                        print("无法从语音识别服务中获取结果; {0}".format(e))
                        await websocket.send_json({"msg": "无法从语音识别服务中获取结果; {0}".format(e),"type":"text", "code": 2001})
                elif msg_type == "file":
                    file_url = body.get("url")  # 假设文件链接作为"url"字段发送
                    # 处理文件链接
                else:
                    # 未知消息类型
                    await websocket.send_json({"msg": "Unknown message type","type":"text", "code": 2001})


                agent_id="666666"
                role_prompt = get_Miku_prompt()
                kg_data = query_kg_message(user_id, agent_id)
                history_messages = find_similar_messages(question,user_id)
                context = ''




                if 'AI资讯' in question or 'AI新闻' in question:
                    context = get_AInews()

                if '帮我查找' in question or '搜索一下' in question or '查一下' in question or '价格' in question or '新闻' in question or '资讯' in question:
                    context = get_bing_search()


                if '提醒我' in question or '记得提醒我' in question or '记得提前告诉我一下' in question:
                    # 创建调度器实例
                    data = get_agent_task(question, user_id)
                    # 创建后台调度器

                    print("调度器已启动，开始接受任务输入。")

                    date = data["date"]  # 2024-03-23
                    time = data["time"]  # 12:01:01
                    event = data["event"]  # 事件
                    run_date = date + " " + time
                    # 添加定时任务，设置任务在2023年7月31日10点30分触发
                    scheduler.add_job(job, 'date', run_date=run_date, args=[websocket,event, user_id])
                    # 启动调度器

                    print(f'时间:{date} {time},事件:{event}, 用户{user_id}')
                    print("定时任务已创建。")
                    response = f"【用户】{user_id}\n【时间】{date} {time}\n【事件】{event}\n已为您创建日程！"

                    content = {"msg": "定时任务已创建", "data": response,"type":"text", "code": 2000}

                    await websocket.send_json(content)


                else:
                    promptTemplate = get_promptTemplate(role_prompt, username, kg_data, history_messages, context)
                    Messages = get_latest_messages(user_id, agent_id, promptTemplate, question)

                    ##非流式输出
                    # response = llm_ZhipuAI(Messages)
                    # st.write(response)

                    ##流式输出
                    response = ""
                    # for output in st.write_stream(llm_stream_ZhipuAI(Messages)):
                    #     response += output  # 在后台收集输出
                    res_content = client.chat.completions.create(
                        model="glm-4",  # 填写需要调用的模型名称
                        messages=Messages,
                        stream=True,
                    )
                    for chunk in res_content:
                        response += chunk.choices[0].delta.content
                        await websocket.send_json({"msg": "200","data":chunk.choices[0].delta.content, "type": "text", "code": 2001})

                    # 写入数据库
                    audio_data = generate_tts(response)

                    # 检查响应是否成功
                    if response is not None:
                        # 将音频数据包装在BytesIO对象中
                        sound = BytesIO(audio_data)

                        # 创建一个包含类型信息的JSON对象
                        message = {
                            "type": "audio",  # 表示这是一个音频消息
                            "msg":"获取成功",
                            "data": sound.read() ,# 音频数据的二进制内容
                            "code": 2000
                        }

                        # 将包含类型信息的JSON对象作为文本消息发送
                        await websocket.send_json(message)
                    else:
                        # 发送错误消息
                        await websocket.send_json({"msg": "TTS生成失败", "type": "text", "code": 2001})

                result, total_msgs = query_msg_counts(user_id, agent_id)
                if result != None:
                    insert_chat_message(result, user_id, agent_id, timestamp=None)
                    if total_msgs % 60 == 0:
                        kg = get_Knowledge_Graph(kg_data, context)
                        insert_or_update_kg_message(user_id, agent_id, kg)
                insert_message(user_id, agent_id, "user", question)
                insert_message(user_id, agent_id, "assistant", response)
            except json.JSONDecodeError as e:
                await websocket.send_json({"msg": "Invalid JSON", "error": str(e), "type": "error", "code": 2001})
            except Exception as e:
                await websocket.send_json({"msg": "An error occurred", "error": str(e), "type": "error", "code": 2001})
    except WebSocketDisconnect:
        # 客户端断开了连接，执行清理操作
        print("WebSocket connection closed by the client.")
    except Exception as e:
        # 发生异常，处理异常并关闭连接
        print(f"An error occurred: {e}")
    finally:
        # 无论发生什么，最终都关闭WebSocket连接
        await websocket.close()








@router.post('/insert_user')
async def insert_user(body=Body(None)):
    user_id=body["user_id"]
    user_name = body["user_name"]
    create_at=datetime.now()

    # Connect to the database
    connection = pymysql.connect(
        host='',
        user='ai_agent',
        password='wEWD4Tipe6ZLm8rW',
        database='ai_agent',
        cursorclass=pymysql.cursors.DictCursor
    )
    try:
        with connection.cursor() as cursor:
            # Create a new record
            sql = "INSERT INTO users (user_id, create_at, user_name) VALUES (%s, %s, %s)"
            cursor.execute(sql, (user_id, create_at, user_name))

        # Commit changes
        connection.commit()

    except Exception as e:
        print(f"Error inserting user: {e}")
        # Rollback in case of any error
        connection.rollback()
    finally:
        connection.close()

@router.post('/query_user')
async def query_user(body=Body(None)):
    user_id=body["user_id"]

    # Connect to the database
    connection = pymysql.connect(
        host='',
        user='ai_agent',
        password='wEWD4Tipe6ZLm8rW',
        database='ai_agent',
        cursorclass=pymysql.cursors.DictCursor
    )
    try:
        with connection.cursor() as cursor:
            # Read a single record
            sql = "SELECT user_id, user_name FROM users WHERE user_id = %s"
            cursor.execute(sql, (user_id,))
            result = cursor.fetchone()
            return result

    except Exception as e:
        print(f"Error querying user: {e}")
    finally:
        connection.close()

from pymysql.cursors import SSCursor
# Close the connection

@router.post('/get_messages')
async def get_messages(body=Body(None)):
    # 使用示例
    user_id=body["user_id"]

    db_config = {
        'host': '',
        'user': 'ai_agent',
        'password': 'wEWD4Tipe6ZLm8rW',
        'vector_db': 'ai_agent',
        'charset': 'utf8mb4',
        'cursorclass': SSCursor
    }

    # 数据库连接配置
    connection = pymysql.connect(**db_config)
    agent_id="666666"
    try:
        with connection.cursor() as cursor:
            # 构建SQL查询语句
            sql = """
            SELECT datetime, sender, content
            FROM messages
            WHERE user_id = %s AND agent_id = %s
            ORDER BY datetime DESC, sender ASC 
            """

            # 执行SQL查询
            cursor.execute(sql, (user_id, agent_id))

            # 获取查询结果
            messages = cursor.fetchall()

        return JSONResponse(content={"msg": "获取成功", "data": messages, "code": 2000})


    except Exception as e:

        # Convert the exception to a string for JSON serialization

        return JSONResponse(content={"msg": "获取失败", "data": str(e), "code": 2001})














if __name__ == "__main__":
    import uvicorn

    app.include_router(router)
    uvicorn.run(app,host="0.0.0.0", port=8009)
#uvicorn main:app --host 0.0.0.0 --port 8009 --reload

# pip freeze > requirements.txt
